// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/array/builder_union.h>

#include <cstddef>
#include <utility>

#include <nebula/core/buffer.h>
#include <turbo/base/checked_cast.h>
#include <turbo/log/logging.h>

namespace nebula {

    turbo::Status BasicUnionBuilder::finish_internal(std::shared_ptr<ArrayData> *out) {
        int64_t length = types_builder_.length();

        std::shared_ptr<Buffer> types;
        TURBO_RETURN_NOT_OK(types_builder_.finish(&types));

        std::vector<std::shared_ptr<ArrayData>> child_data(children_.size());
        for (size_t i = 0; i < children_.size(); ++i) {
            TURBO_RETURN_NOT_OK(children_[i]->finish_internal(&child_data[i]));
        }

        *out = ArrayData::create(type(), length, {nullptr, types}, /*null_count=*/0);
        (*out)->child_data = std::move(child_data);
        return turbo::OkStatus();
    }

    turbo::Status DenseUnionBuilder::append_array_slice(const ArraySpan &array, const int64_t offset,
                                                      const int64_t length) {
        const int8_t *type_codes = array.get_values<int8_t>(1);
        const int32_t *offsets = array.get_values<int32_t>(2);
        for (int64_t row = offset; row < offset + length; row++) {
            const int8_t type_code = type_codes[row];
            const int child_id = type_id_to_child_id_[type_code];
            const int32_t union_offset = offsets[row];
            TURBO_RETURN_NOT_OK(append(type_code));
            TURBO_RETURN_NOT_OK(type_id_to_children_[type_code]->append_array_slice(
                    array.child_data[child_id], union_offset, /*length=*/1));
        }
        return turbo::OkStatus();
    }

    turbo::Status DenseUnionBuilder::finish_internal(std::shared_ptr<ArrayData> *out) {
        TURBO_RETURN_NOT_OK(BasicUnionBuilder::finish_internal(out));
        (*out)->buffers.resize(3);
        TURBO_RETURN_NOT_OK(offsets_builder_.finish(&(*out)->buffers[2]));
        return turbo::OkStatus();
    }

    BasicUnionBuilder::BasicUnionBuilder(
            MemoryPool *pool, int64_t alignment,
            const std::vector<std::shared_ptr<ArrayBuilder>> &children,
            const std::shared_ptr<DataType> &type)
            : ArrayBuilder(pool, alignment),
              child_fields_(children.size()),
              types_builder_(pool, alignment) {
        const auto &union_type = turbo::checked_cast<const UnionType &>(*type);
        mode_ = union_type.mode();

                DKCHECK_EQ(children.size(), union_type.type_codes().size());

        type_codes_ = union_type.type_codes();
        children_ = children;

        type_id_to_child_id_.resize(union_type.max_type_code() + 1, -1);
        type_id_to_children_.resize(union_type.max_type_code() + 1, nullptr);
                DKCHECK_LE(
                type_id_to_children_.size() - 1,
                static_cast<decltype(type_id_to_children_)::size_type>(UnionType::kMaxTypeCode));

        for (size_t i = 0; i < children.size(); ++i) {
            child_fields_[i] = union_type.field(static_cast<int>(i));

            auto type_id = union_type.type_codes()[i];
            type_id_to_child_id_[type_id] = static_cast<int>(i);
            type_id_to_children_[type_id] = children[i].get();
        }
    }

    int8_t BasicUnionBuilder::AppendChild(const std::shared_ptr<ArrayBuilder> &new_child,
                                          const std::string &field_name) {
        children_.push_back(new_child);
        auto new_type_id = NextTypeId();

        type_id_to_child_id_[new_type_id] = static_cast<int>(children_.size() - 1);
        type_id_to_children_[new_type_id] = new_child.get();
        child_fields_.push_back(field(field_name, nullptr));
        type_codes_.push_back(static_cast<int8_t>(new_type_id));

        return new_type_id;
    }

    std::shared_ptr<DataType> BasicUnionBuilder::type() const {
        std::vector<std::shared_ptr<Field>> child_fields(child_fields_.size());
        for (size_t i = 0; i < child_fields.size(); ++i) {
            child_fields[i] = child_fields_[i]->with_type(children_[i]->type());
        }
        return mode_ == UnionMode::SPARSE ? sparse_union(std::move(child_fields), type_codes_)
                                          : dense_union(std::move(child_fields), type_codes_);
    }

    int8_t BasicUnionBuilder::NextTypeId() {
        // Find type_id such that type_id_to_children_[type_id] == nullptr
        // and use that for the new child. Start searching at dense_type_id_
        // since type_id_to_children_ is densely packed up at least up to dense_type_id_
        for (; static_cast<size_t>(dense_type_id_) < type_id_to_children_.size();
               ++dense_type_id_) {
            if (type_id_to_children_[dense_type_id_] == nullptr) {
                return dense_type_id_++;
            }
        }

                DKCHECK_LT(
                type_id_to_children_.size(),
                static_cast<decltype(type_id_to_children_)::size_type>(UnionType::kMaxTypeCode));

        // type_id_to_children_ is already densely packed, so just append the new child
        type_id_to_child_id_.resize(type_id_to_child_id_.size() + 1);
        type_id_to_children_.resize(type_id_to_children_.size() + 1);
        return dense_type_id_++;
    }

    turbo::Status SparseUnionBuilder::append_array_slice(const ArraySpan &array, const int64_t offset,
                                                       const int64_t length) {
        for (size_t i = 0; i < type_codes_.size(); i++) {
            TURBO_RETURN_NOT_OK(type_id_to_children_[type_codes_[i]]->append_array_slice(
                    array.child_data[i], array.offset + offset, length));
        }
        const int8_t *type_codes = array.get_values<int8_t>(1);
        TURBO_RETURN_NOT_OK(types_builder_.append(type_codes + offset, length));
        return turbo::OkStatus();
    }

}  // namespace nebula
